import os
import numpy as np
import cv2
import torch
from torch.nn import functional as F


def flow_warp(x, flo):
    """
    inverse warp an image/tensor (im2) back to im1, according to the optical flow

    x: [B, C, H, W] (im2)
    flo: [B, 2, H, W] flow

    """
    B, C, H, W = x.size()
    # mesh grid
    xx = torch.arange(0, W).view(1, -1).repeat(H, 1).to(x.device)
    yy = torch.arange(0, H).view(-1, 1).repeat(1, W).to(x.device)
    xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1)
    yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1)
    grid = torch.cat((xx, yy), 1).float()

    vgrid = grid + flo

    # scale grid to [-1,1]
    vgrid = torch.stack([
        2.0 * vgrid[:, 0, :, :] / max(W - 1, 1) - 1.0,
        2.0 * vgrid[:, 1, :, :] / max(H - 1, 1) - 1.0
    ],
                        dim=1)

    vgrid = vgrid.permute(0, 2, 3, 1)
    output = F.grid_sample(x, vgrid, mode='nearest', padding_mode='border')

    return output


def warping(mask_fn, flow_fn, pred_fn):
    mask = cv2.imread(mask_fn, cv2.IMREAD_GRAYSCALE).astype(np.float)
    flow = cv2.imread(flow_fn).astype(np.float)[:, :, :2]
    # new mask
    img = torch.Tensor(mask).unsqueeze(dim=0).unsqueeze(dim=1)
    flow = torch.Tensor(flow).permute(2, 0, 1).contiguous().unsqueeze(dim=0)

    new_mask = flow_warp(img, flow)
    cv2.imwrite(pred_fn, new_mask)
    print(pred_fn)


def main():
    gt_base = '/shared/xiali/datasets/DAVIS/Annotations/2018'
    fl_base = '/shared/xiali/outputs/semi-flow/davis_flow_val'
    pd_base = '/shared/xiali/outputs/semi-flow/davis_mask_val'
    if not os.isfolder(pd_base):
        os.makedirs(pd_base)

    for folder in os.listdir(gt_base):
        folder_path = os.path.join(gt_base, folder)
        if not os.isfolder(folder_path):
            os.makedirs(folder_path)
        names = os.listdir(folder_path)
        for i in range(len(names) - 1):
            mask_fn = os.path.join(gt_base, names[i])
            flow_fn = os.path.join(fl_base, names[i].split('.')[0] + '.flo')
            pred_fn = os.path.join(pd_base, names[i])
            warping(mask_fn, flow_fn, pred_fn)


if __name__ == "__main__":
    main()
